import re
import requests
import os
import concurrent.futures


def download_ts_files(m3u8_url, output_dir='ts_files'):
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 检查m3u8文件是否存在
    m3u8_path = os.path.join(output_dir, 'index.m3u8')
    if os.path.exists(m3u8_path):
        print('m3u8文件已存在，直接读取')
        with open(m3u8_path, 'r') as f:
            m3u8_content = f.read()
    else:
        # 下载m3u8文件
        print('正在下载m3u8文件')
        response = requests.get(m3u8_url)
        m3u8_content = response.text
        with open(m3u8_path, 'w') as f:
            f.write(m3u8_content)
            
    # 解析m3u8文件，获取.ts文件URL
    ts_prefix = '/'.join(m3u8_url.split('/')[:-1])
    ts_urls = [f'{ts_prefix}/{ts}' for ts in re.findall(r'[^\s]+\.ts', m3u8_content)]

    # 使用线程池下载.ts文件
    with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
        future_to_url = {executor.submit(download_single_ts, ts_url, output_dir): ts_url for ts_url in ts_urls}
        for future in concurrent.futures.as_completed(future_to_url):
            try:
                future.result()
            except Exception as exc:
                print(f'下载过程中发生错误: {exc}')

            # 显示下载进度
            downloaded_count = len([f for f in os.listdir(output_dir) if f.endswith('.ts') and not f.startswith('.')])
            total_count = len(ts_urls)
            print(f'下载进度: {downloaded_count}/{total_count} ({downloaded_count/total_count:.2%})')

    print(f'所有.ts文件已下载到 {output_dir} 目录')


def download_single_ts(ts_url, output_dir):
    file_name = ts_url.split('/')[-1]
    file_path = os.path.join(output_dir, file_name)

    # 检查文件是否存在
    if os.path.exists(file_path):
        print(f'文件已存在，跳过下载: {file_name}')
        return

    # 下载文件，最多重试3次
    max_retries = 3
    for attempt in range(max_retries):
        try:
            print(f'正在下载: {file_name} (尝试 {attempt + 1}/{max_retries})')
            response = requests.get(ts_url)
            with open(file_path, 'wb') as f:
                f.write(response.content)
            break
        except Exception as e:
            print(f'下载失败: {file_name}, 错误: {str(e)}')
            if attempt == max_retries - 1:
                print(f'文件 {file_name} 下载失败，已达到最大重试次数')


if __name__ == '__main__':
    # 示例用法
    # m3u8_url = 'https://vip.lz-cdn11.com/20220717/9411_ca9387e9/1000k/hls/mixed.m3u8'
    # download_ts_files(m3u8_url, '/Users/yawn/Geil/m3u8/ssbx-all')

    # m3u8_url = 'https://vip.lz-cdn11.com/20220425/4284_8808bfcc/1000k/hls/mixed.m3u8'
    # save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220425-1'
    # download_ts_files(m3u8_url, save_dir)

    # m3u8_url = 'https://vip.lz-cdn11.com/20220425/4291_44bb1246/1000k/hls/mixed.m3u8'
    # save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220425-2'
    # download_ts_files(m3u8_url, save_dir)

    # m3u8_url = 'https://vip.lz-cdn11.com/20220502/5233_0b33d1b6/1000k/hls/mixed.m3u8'
    # save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220502'
    # download_ts_files(m3u8_url, save_dir)

    # m3u8_url = 'https://vip.lz-cdn11.com/20220509/5811_ce8b8bdf/1000k/hls/mixed.m3u8'
    # save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220509'
    # download_ts_files(m3u8_url, save_dir)

    # m3u8_url = 'https://vip.lz-cdn11.com/20220522/7302_22fac913/1000k/hls/mixed.m3u8'
    # save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220522'
    # download_ts_files(m3u8_url, save_dir)

    # m3u8_url = 'https://vip.lz-cdn11.com/20220605/8042_e4f6e440/1000k/hls/mixed.m3u8'
    # save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220605'
    # download_ts_files(m3u8_url, save_dir)

    # m3u8_url = 'https://vip.lz-cdn11.com/20220509/5811_ce8b8bdf/1000k/hls/mixed.m3u8'
    # save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220509'
    # download_ts_files(m3u8_url, save_dir)

    # m3u8_url = 'https://vip.lz-cdn11.com/20220605/8041_2172e1b4/1000k/hls/mixed.m3u8'
    # save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220605'
    # download_ts_files(m3u8_url, save_dir)

    m3u8_url = 'https://vip.lz-cdn11.com/20220612/8407_42b3bc8c/1000k/hls/mixed.m3u8'
    save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220612'
    download_ts_files(m3u8_url, save_dir)

    m3u8_url = 'https://vip.lz-cdn11.com/20220619/8771_379ea2c2/1000k/hls/mixed.m3u8'
    save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220619'
    download_ts_files(m3u8_url, save_dir)

    m3u8_url = 'https://vip.lz-cdn12.com/20220627/336_b1456a1a/1200k/hls/mixed.m3u8'
    save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220627'
    download_ts_files(m3u8_url, save_dir)

    m3u8_url = 'https://vip.lz-cdn11.com/20220704/9095_129f9c41/1000k/hls/mixed.m3u8'
    save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220704'
    download_ts_files(m3u8_url, save_dir)

    m3u8_url = 'https://vip.lz-cdn11.com/20220711/9266_6b9ba672/1000k/hls/mixed.m3u8'
    save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220711-1'
    download_ts_files(m3u8_url, save_dir)

    m3u8_url = 'https://vip.lz-cdn11.com/20220711/9267_88b0f00a/1000k/hls/mixed.m3u8'
    save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220711-2'
    download_ts_files(m3u8_url, save_dir)

    m3u8_url = 'https://vip.lz-cdn11.com/20220715/9383_5798ad7c/1000k/hls/mixed.m3u8'
    save_dir = '/Users/yawn/Geil/m3u8/ssbx-20220715'
    download_ts_files(m3u8_url, save_dir)